# %%
from conn import clip_annotated_record, check_frame_label_col
from datetime import datetime, timedelta
from tqdm import tqdm
from collections import defaultdict

# %%

groups = check_frame_label_col.aggregate(
    [
        {
            "$group": {
                "_id": {"anno_type": "$annotate_type"},
                "clip_ids": {"$addToSet": "$clip_id"},
            }
        }
    ]
).to_list()
print("group success")

# %%

for group in groups:
    anno_type = group["_id"]["anno_type"]
    for clip_id in group["clip_ids"]:
        frames = check_frame_label_col.find(
            {"annotate_type": anno_type, "clip_id": clip_id},
            {"created_at": 1, "frame_ids": 1},
        ).sort("created_at", -1)
        for frame in frames:
            assert len(frame["frame_ids"]) == 1

# %%
data = check_frame_label_col.aggregate(
    [
        {
            "$addFields": {
                "frame_cnt": {"$size": "$frame_ids"},
            }
        },
        {"$match": {"frame_cnt": {"$gt": 1}}},
        {"$project": {"_id": 1, "annotate_type": 1, "clip_id": 1, "frame_cnt": 1}},
    ]
).to_list()
print(len(data))

# %%
print(len([item for item in data if item.get("annotate_type") is None]))
# for item in data:
#     if item['annotate_type']:
#         print(item)

# %%
print(len(groups))

for group in groups:
    anno_type = group["_id"]["anno_type"]
    print(anno_type, len(group["clip_ids"]))

# %%

repeat_cnt = defaultdict(int)
for group in groups:
    anno_type = group["_id"]["anno_type"]
    if anno_type is None:
        continue
    print(anno_type, len(group["clip_ids"]))
    for clip_id in tqdm(group["clip_ids"]):
        frames = check_frame_label_col.find(
            {"annotate_type": anno_type, "clip_id": clip_id},
            {"created_at": 1, "frame_ids": 1},
        ).sort("created_at", -1)
        frame_id_set = set()
        for frame in frames:
            target_frame_id = frame["frame_ids"][0]
            if target_frame_id in frame_id_set:
                repeat_cnt[anno_type] += 1
            else:
                frame_id_set.add(target_frame_id)
    break
print(repeat_cnt)
# %%
for group in groups:
    print(group["_id"]["annotate_type"])

# %%
# 1 删除重复帧
# 2 更新 clip 状态和验收完成时间
# 3 后端统计 finished ac clip 数判断 ac task 是否完成，移除定时任务

# %%
repeat_cnt = defaultdict(int)
for group in groups:
    anno_type = group["_id"]["anno_type"]
    if anno_type is None:
        continue
    for clip_id in tqdm(group["clip_ids"]):
        frame_cnt = check_frame_label_col.count_documents(
            {"annotate_type": anno_type, "clip_id": clip_id},
        )
        assert frame_cnt > 0, "frame cnt must > 0"
        latest_row = check_frame_label_col.find(
            {"annotate_type": anno_type, "clip_id": clip_id}, {"created_at": 1}
        ).sort({"created_at": -1})[0]
        ac_finished_at = latest_row["created_at"]


# %%
# 以 anno-type 为例，看下实际有多少帧

frame_set = set()
latest_time = None
anno_type = 31
total = check_frame_label_col.count_documents({"annotate_type": anno_type})
for frame in tqdm(
    check_frame_label_col.find(
        {"annotate_type": anno_type}, {"frame_ids": 1, "created_at": 1}
    ).sort({"created_at": -1}),
    total=total,
):
    assert len(frame["frame_ids"]) == 1
    frame_set.add(frame["frame_ids"][0])
    if latest_time is None:
        latest_time = frame["created_at"]
print(len(frame_set), total, latest_time)


# %%
# 以 anno-type 为例，找到重复帧并删除
frame_set = set()
latest_time = None
anno_type = 31
total = check_frame_label_col.count_documents({"annotate_type": anno_type})
batch = 1000
del_ids = []
del_cnt = 0
for frame in tqdm(
    check_frame_label_col.find(
        {"annotate_type": anno_type}, {"frame_ids": 1, "created_at": 1, "clip_id": 1}
    ).sort({"created_at": -1}),
    total=total,
):
    assert len(frame["frame_ids"]) == 1
    target_frame_id = frame["frame_ids"][0]
    if target_frame_id in frame_set:
        # print(frame["clip_id"], target_frame_id, frame["_id"])
        # check_frame_label_col.delete_one({"_id": frame["_id"]})
        del_ids.append(frame["_id"])
        if len(del_ids) >= batch:
            check_frame_label_col.delete_many({"_id": {"$in": del_ids}})
            del_cnt += len(del_ids)
            del_ids = []
    else:
        frame_set.add(target_frame_id)
if len(del_ids) > 0:
    check_frame_label_col.delete_many({"_id": {"$in": del_ids}})
    del_cnt += len(del_ids)
    del_ids = []
print(del_cnt)


# %%
# 更新验收状态和验收完成时间
for group in groups:
    anno_type = group["_id"]["anno_type"]
    if anno_type is None:
        continue
    print(anno_type, len(group["clip_ids"]))
    for clip_id in tqdm(group["clip_ids"]):
        rows = (
            check_frame_label_col.find(
                {"annotate_type": anno_type, "clip_id": clip_id}, {"created_at": 1}
            )
            .sort({"created_at": -1})
            .limit(1)
            .to_list()
        )
        assert len(rows) == 1
        latest_row = rows[0]
        rows = clip_annotated_record.find(
            {"annotate_type": anno_type, "clip_id": clip_id},
            {"_id": 1},
        ).to_list()
        if len(rows) != 1:
            print(len(rows), rows, clip_id)
        target_row_id = rows[0]["_id"]
        clip_annotated_record.update_one(
            {"_id": target_row_id},
            {
                "$set": {
                    "accept_status": 20,  # 已完成
                    # 使用最近一次提交时间
                    "accept_finished_at": latest_row["created_at"],
                }
            },
        )
